# 稳定版下载脚本 - 支持文件存在自动跳过
from huggingface_hub import snapshot_download
import os
import time
import requests

# 设置环境变量
os.environ['HF_HUB_CACHE'] = '.././models/hf_cache'
os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = '300'  # 5分钟超时

def check_network():
    """检查网络连接"""
    try:
        response = requests.get('https://huggingface.co', timeout=10)
        return response.status_code == 200
    except:
        return False

def is_model_downloaded(local_dir):
    """判断模型文件夹是否存在且非空"""
    return os.path.exists(local_dir) and os.path.isdir(local_dir) and len(os.listdir(local_dir)) > 0

def download_with_patience(repo_id, local_dir, max_retries=2):
    """耐心下载，单线程避免卡顿"""
    if is_model_downloaded(local_dir):
        print(f"✅ {repo_id} 已存在，跳过下载。")
        return True

    for attempt in range(max_retries):
        try:
            print(f"🔄 尝试下载 {repo_id} (第 {attempt + 1} 次)")

            # 检查网络
            if not check_network():
                print("❌ 网络连接异常，请检查网络")
                return False

            # 使用单线程下载，避免卡顿
            snapshot_download(
                repo_id=repo_id,
                local_dir=local_dir,
                resume_download=True,
                max_workers=1,  # 单线程下载
            )
            print(f"✅ {repo_id} 下载成功！")
            return True

        except Exception as e:
            error_msg = str(e)
            print(f"❌ 下载失败: {error_msg[:100]}...")

            if "timeout" in error_msg.lower() or "connection" in error_msg.lower():
                print("🔄 网络超时，准备重试...")
                if attempt < max_retries - 1:
                    print("⏳ 等待 15 秒后重试...")
                    time.sleep(15)
                continue
            else:
                print(f"💔 严重错误，停止重试")
                return False

    return False

# 主程序
print("🌐 检查网络连接...")
if not check_network():
    print("❌ 无法连接到 Hugging Face，请检查网络")
    exit(1)

print("✅ 网络连接正常")
print("🚀 开始下载模型（单线程模式，更稳定）...")

# 创建目录
os.makedirs(".././models", exist_ok=True)

# 需要下载的模型列表
model_list = [
    # 翻译模型
    ("Helsinki-NLP/opus-mt-zh-en", ".././models/opus-mt-zh-en"),
    # 文本生成模型
    ("succinctly/text2image-prompt-generator", ".././models/text2image-prompt-generator"),
    # CLIP 模型
    ("laion/CLIP-ViT-H-14-laion2B-s32B-b79K", ".././models/clip-vit-h-14"),
    # BLIP-Large 模型
    ("Salesforce/blip-image-captioning-large", ".././models/blip-large"),
]

results = []
for repo_id, local_dir in model_list:
    print(f"\n📥 检查/下载 {repo_id} ...")
    success = download_with_patience(repo_id=repo_id, local_dir=local_dir)
    results.append(success)

# 结果总结
print("\n" + "=" * 50)
if all(results):
    print("🎉 所有模型准备完成！")
else:
    print("❌ 部分模型未准备好，可重新运行脚本续传")
